feat: Add SLURM scheduler support for ORB#246
Conversation
fgogolli
changed the title
wettowelreactor
force-pushed
the
feat/slurm-scheduler
branch
from
e5f3fc0 to
dff63d9
Compare
There was a problem hiding this comment.
This PR adds SLURM as a third scheduler alongside default and hostfactory. The core addition - the src/orb/infrastructure/scheduler/slurm/ tree, registration wiring, and the new generate_scheduler_templates port extension - follows the same structural pattern as hostfactory. The SchedulerPort extension has a backward-compatible None default, registration uses the existing dict dispatch, and SlurmFieldMapper inherits SchedulerFieldMapper cleanly.
The dominant action item before merge is a rebase onto current main (dd4c3903). The PR was cut before #253 merged, so a large portion of the diff reflects pre-#253 baseline drift rather than intentional changes. After rebase and removal of the committed artefact described below, the SLURM-specific scope shrinks substantially.
Blockers that need active work (not resolved by rebase alone):
coverage-unit.xml (42,326 lines) is committed at the repo root and contains the absolute path /Users/rouc/dev/open-resource-broker/src/orb - a contributor's local filesystem layout exposed in a public repository. The file matches neither coverage.xml nor coverage/ in .gitignore. Remove it with git rm coverage-unit.xml, add coverage-unit.xml to .gitignore, and amend before merge.
slurmmock.py sits at the repo root, importable by any code that appends . to sys.path, and absent from .gitignore. Move it to tests/integration/slurm/slurmmock.py where it already has a natural home alongside test_slurm_lifecycle.py.
tests/unit/infrastructure/scheduler/conftest.py imports from tests.onaws.plugin_io_schemas. After rebase onto main, that path does not exist - #253 moved the module to tests/providers/aws/live/plugin_io_schemas.py. The stale import kills pytest collection for the entire tests/unit/infrastructure/scheduler/ subtree, which means default and hostfactory scheduler tests also fail. Update the import path and add the three SLURM-specific schema aliases to the new location (or to a SLURM-specific file there).
_resolve_template_for_nodes in slurm_strategy.py returns TemplateDTO(template_id="default") unconditionally. Every ResumeProgram invocation passes through this method, so any multi-partition cluster - one with a gpu or compute partition mapped to a distinct template - silently provisions the wrong instance type, AMI, and capacity config. The integration tests bypass this path entirely via MockSlurmAppService, so nothing currently catches it. Either raise NotImplementedError with a clear message so operators see the gap immediately, or implement the partition-to-template lookup; shipping a silent wrong-template fallback is harder to diagnose in production than a clean error.
src/orb/interface/request_command_handlers.py imports SlurmNodeMapper directly from infrastructure/scheduler/slurm/node_mapper at two call sites (lines 116 and 215). The interface layer must not reach into infrastructure. The handler already holds a SchedulerPort reference (line 99); add expand_node_range(node_spec: str) -> list[str] to SchedulerPort and route through that. The method is pure string logic with no I/O, so it can live on BaseSchedulerStrategy as a static method - SLURM has no monopoly on bracket notation.
node_bootstrap.py validates node_name and slurmctld_host against _NAME_RE before embedding them in the cloud-init shell heredoc, but slurm_conf_path is interpolated bare with no validation. Add an allowlist regex (^[/a-zA-Z0-9._-]+$) and wrap the interpolation with shlex.quote. The generated script runs as root on a fresh instance, so the blast radius of an injected payload is full instance compromise.
resumeProgram.sh and suspendProgram.sh construct the JSON payload with bash string substitution (${NODE_LIST// /", "}). When SLURM passes a bracket-expanded hostlist such as compute-[001-003], the substitution produces a single-element array containing the literal bracket string rather than the three expanded node names. scontrol show hostnames is already called for NUM_NODES - capture that output and feed the expanded names into jq --argjson or jq -R to build the array safely.
Items for this PR or a fast-follow:
_sync_machine_tags_to_provider in request_creation_handlers.py guards with str(m.machine_id).startswith("i-"). That guard encodes AWS EC2 instance-ID format in the provider-agnostic application layer. For now it is only silently a no-op for non-AWS providers, not a regression, but the correct seam is ProviderCapabilityPort.supports_operation(TAG_INSTANCES) or moving the guard into the AWS provider strategy. If this change ships in this PR, the i- check needs to go.
SLURM_ORB_RESTD_URL accepts any http:// or https:// URL with no hostname allowlist. As an operator-configured value this is a low-severity SSRF surface, but the deployment guide should document that the variable must point to a known-trusted slurmrestd host and that HTTPS with a valid certificate is the recommended default, not the plain-HTTP local example.
OperationType.TAG_INSTANCES in domain/base/operations.py duplicates ProviderOperationType.TAG_INSTANCES in providers/base/strategy/provider_strategy.py; both resolve to the string "tag_instances". This predates the PR but the new SLURM tag-sync path adds a third call site that relies on the accidental value equality. Consolidate on one enum.
node_bootstrap.py, rest_client.py, and cli_adapter.py use logging.getLogger(__name__) rather than the constructor-injected LoggingPort pattern used throughout the AWS provider. They will not participate in structured log routing or correlation-ID propagation. Align with the established pattern.
The SLURM schemas in tests/providers/aws/live/plugin_io_schemas.py (the correct post-rebase path) currently alias directly to the default scheduler schemas. The parametric scheduler_config fixture includes SLURM in the matrix, giving the appearance of three-scheduler coverage while actually asserting nothing SLURM-specific. Add dedicated SLURM schemas that validate partition_name, node_list, and other SLURM envelope fields.
handle_resume_request, handle_suspend_request, SlurmCliAdapter, SlurmRestClient, and generate_scheduler_templates (including _parse_slurm_conf and _match_instance_type) have no direct test coverage - the lifecycle tests route through MockSlurmAppService and never touch the real strategy methods. Replace at least the happy-path lifecycle test with real-strategy tests that mock at the HTTP or subprocess boundary rather than above the strategy class.
| HostFactorySchedulerStrategy, | ||
| ) | ||
| from orb.infrastructure.scheduler.slurm.slurm_strategy import SlurmSchedulerStrategy | ||
| from tests.onaws.plugin_io_schemas import ( |
There was a problem hiding this comment.
After rebase onto main, this import path won't resolve. #253 moved the module from tests/onaws/ to tests/providers/aws/live/plugin_io_schemas.py. Update the import here and move the three SLURM-specific schema aliases (expected_*_schema_slurm) to that file. Without this fix, pytest fails to collect the entire tests/unit/infrastructure/scheduler/ subtree — taking default and hostfactory scheduler tests down with it.
| @@ -0,0 +1,219 @@ | |||
| #!/usr/bin/env python | |||
| """Mock slurmrestd server for local development and testing. | |||
There was a problem hiding this comment.
slurmmock.py belongs under tests/integration/slurm/ where test_slurm_lifecycle.py already lives. At the repo root it sits alongside production source, will be picked up by any glob that sweeps *.py from ., and has no .gitignore entry. Move it before merge.
| } | ||
| node_names = getattr(args, "nodes", None) | ||
| if node_names: | ||
| from orb.infrastructure.scheduler.slurm.node_mapper import SlurmNodeMapper |
There was a problem hiding this comment.
SlurmNodeMapper is imported directly from infrastructure/scheduler/slurm/ at lines 116 and 215. The handler already holds a SchedulerPort reference (line 99); add expand_node_range(node_spec: str) -> list[str] to SchedulerPort and route through that. The expansion logic is pure string manipulation with no I/O or SLURM-specific state — it can live on BaseSchedulerStrategy as a static method.
| hostnamectl set-hostname {node_name} | ||
|
|
||
| # Ensure slurm.conf has correct SlurmctldHost | ||
| sed -i 's/^SlurmctldHost=.*/SlurmctldHost={slurmctld_host}/' {slurm_conf_path} |
There was a problem hiding this comment.
slurm_conf_path is interpolated bare into the cloud-init shell heredoc. node_name and slurmctld_host are guarded by _NAME_RE; this parameter is not. The generated script runs as root on a fresh instance. Add an allowlist regex (^[/a-zA-Z0-9._-]+$) and wrap the interpolation site with shlex.quote before merge.
|
|
||
| # --- Single batch request to ORB --- | ||
| if [ "${ORB_MODE}" = "api" ]; then | ||
| PAYLOAD="{\"node_names\": [\"${NODE_LIST// /\", \"}\"], \"request_type\": \"provision\"}" |
There was a problem hiding this comment.
When SLURM passes a bracket-expanded hostlist (compute-[001-003]), the space-substitution idiom produces ["compute-[001-003]"] — a single-element array, not three expanded names. scontrol show hostnames is already called on line 67 for NUM_NODES; capture that output and build the JSON array from it using jq --argjson or by piping each expanded name through jq -R. The same pattern is in suspendProgram.sh.
| } | ||
| ) | ||
|
|
||
| def _resolve_template_for_nodes(self, node_names: list[str]) -> Any: |
There was a problem hiding this comment.
_resolve_template_for_nodes always returns TemplateDTO(template_id="default") regardless of partition. handle_resume_request calls this on every resume, so multi-partition clusters silently provision the wrong template. The integration tests route through MockSlurmAppService and never exercise this path, so there is no current safety net. Raise NotImplementedError until the partition-to-template lookup is implemented, so callers see a clear error rather than a plausible-looking but wrong result.
| instance_tags: dict[str, dict[str, str]] = {} | ||
| for m in machines: | ||
| if m.tags and m.tags.tags and str(m.machine_id).startswith("i-"): | ||
| instance_tags[str(m.machine_id)] = m.tags.tags |
There was a problem hiding this comment.
str(m.machine_id).startswith("i-") encodes the AWS EC2 instance-ID format in the provider-agnostic application layer. Non-AWS providers are silently skipped with no log or error. The correct seam is ProviderCapabilityPort.supports_operation(TAG_INSTANCES), or move this guard into the AWS provider strategy and let the provider no-op cleanly for unrecognised IDs.
| ) | ||
|
|
||
| def _resolve_template_for_nodes(self, node_names: list[str]) -> Any: | ||
| """Resolve which template/partition these node slots belong to. |
There was a problem hiding this comment.
This stub ships without a guard that makes the gap visible at runtime. Every ResumeProgram invocation reaches this method, and every call returns template_id="default" regardless of the SLURM partition. A cluster with distinct gpu or batch partition templates will provision silently against the wrong config. Raise NotImplementedError with a message describing the missing partition-to-template lookup until this is implemented.
| node_names = getattr(args, "nodes", None) | ||
| if node_names: | ||
| from orb.infrastructure.scheduler.slurm.node_mapper import SlurmNodeMapper | ||
|
|
There was a problem hiding this comment.
Direct import of SlurmNodeMapper from infrastructure/scheduler/slurm/. A SchedulerPort reference is already available at line 99. Add expand_node_range(node_spec: str) -> list[str] to the port — the method is stateless string logic with no infrastructure dependency — and call it through the port here and at line 215.
2 participants
PR: Add SLURM Scheduler Integration
What
Adds SLURM workload manager as a new scheduler type for ORB, enabling ORB to act as a cloud resource provider for SLURM clusters via the ResumeProgram/SuspendProgram power management hooks.
Why
ORB currently supports HostFactory (IBM Spectrum Symphony) and a standalone default scheduler. SLURM is the dominant workload manager in HPC and scientific computing. Adding SLURM support enables ORB to provision and manage cloud compute capacity (EC2) on demand for SLURM clusters using its existing provider infrastructure.
The integration follows SLURM's "dynamic slot" model — cloud nodes are treated as fungible capacity slots with fresh instances on each resume cycle. This matches the AWS ParallelCluster model and SchedMD's cloud-bursting recommendations.
How
Architecture: SLURM decides WHEN to scale (scheduling), ORB decides HOW to provision (cloud APIs). Communication flow:
Components added:
SlurmSchedulerStrategy— fullSchedulerPortimplementation with batch resume/suspend handlersSlurmFieldMapper+field_mappings.py— bidirectional translation between SLURM and ORB domain concepts (node states, partition fields)SlurmResponseFormatter— formats ORB domain objects to SLURM-compatible snake_case responsesSlurmNodeMapper— runtime mapping between SLURM node names and cloud instance IDs with hostlist range expansionSlurmNodeBootstrap— post-provisioning node address registration viascontrol updateSlurmRestClient— slurmrestd REST API client for health monitoringSlurmCliAdapter— CLI fallback (sinfo/scontrol) for environments without slurmrestdresumeProgram.sh/suspendProgram.sh— SLURM power hook entry pointsKey design decisions:
Testing
SCHEDULER_CONFIGSextension point)slurmmock.py) for local development without a real clusterpyright: 0 errors, 0 warningsruff: all checks pass